from qdrant_client import QdrantClient, models
from sentence_transformers import SentenceTransformer
import asyncio
from typing import List, Dict, Any

class QdrantService:
    def __init__(self, host: str, port: int, api_key: str = None, model_name: str = "baai/bge-base-zh-v1.5"):
        self.client = QdrantClient(host=host, port=port)
        self.encoder = SentenceTransformer(model_name)  # 文本编码器
    
    async def search(self, 
                   collection_name: str, 
                   query_text: str, 
                   limit: int = 20, 
                   score_threshold: float = 0.01) -> List[Dict[str, Any]]:
        """
        执行分层检索：用子块搜索，返回顶级父块内容
        三步流程：
          1. 将查询文本编码为向量
          2. 检索相关子块
          3. 获取子块对应的顶级父块内容
        """
        # 1. 文本→向量转换
        print(f"开始编码查询文本: {query_text}")
        query_vector = self.encoder.encode(query_text).tolist()
        
        # 2. 检索相关子块（异步执行）
        child_points = await asyncio.to_thread(
            self.client.search,
            collection_name=collection_name,
            query_vector=query_vector,
            query_filter=models.Filter(
                must=[models.FieldCondition(key="type", match=models.MatchValue(value="child"))]
            ),
            limit=limit,
            score_threshold=score_threshold,
            with_payload= ["root_doc_id",]
        )

        # 3.  按关键字模糊查询payload的content字段
        hits, offset = await asyncio.to_thread(
            self.client.scroll,
            collection_name=collection_name,
            scroll_filter=models.Filter(
                must=[
                    models.FieldCondition(
                        key="content", # 要查询的字段名
                        match=models.MatchText(text=query_text) # 使用 *query_text* 模糊查询
                    ),
                    models.FieldCondition(
                        key="type",
                        match=models.MatchValue(value="child")
                    )
                ]
            ),
            limit=limit,
            with_payload=["root_doc_id",]
        )       

        # 3.1 score设置为1
        
        hit_root_doc_ids = set()
        for hit in hits:
            hit_root_doc_ids.add(hit.payload.get("root_doc_id"))
        print(f"模糊查询到的父块ID: {hit_root_doc_ids}")
        len_hit_root_doc_ids = len(hit_root_doc_ids)
        print(f"模糊查询到的父块数: {len_hit_root_doc_ids}")

        # 4. 按root_doc_id分组按最大值计算score
        root_doc_scores = {}
        if len_hit_root_doc_ids < limit:
            for p in child_points:
                root_doc_id = p.payload.get("root_doc_id")
                if root_doc_id:
                    root_doc_scores[root_doc_id] = max(root_doc_scores.get(root_doc_id, 0), p.score)
        
        # 5. 按score排序
        sorted_root_doc_ids = sorted(root_doc_scores.keys(), key=lambda k: root_doc_scores[k], reverse=True)

        # 5.1 过滤出存在的父块ID
        sorted_root_doc_ids = [doc_id for doc_id in sorted_root_doc_ids if doc_id not in hit_root_doc_ids]
       
        # 5.2 将 hit_root_doc_ids排在前面
        sorted_root_doc_ids = list(hit_root_doc_ids) + sorted_root_doc_ids

        # 6. 召回父块内容
        root_points = await self._get_root_points(collection_name, sorted_root_doc_ids)

        # 7. 按sorted_root_doc_ids顺序，累计文档个数，直到超过20个，
        total_docs = 0
        truncated_root_points = []
        for point in root_points:
            if total_docs > 20 :
                break
            total_docs += 1
            truncated_root_points.append(point)


        # 顶级父块内容
        return [{
            "content": point.payload.get("content", ""),
            "score": root_doc_scores.get(point.id, 1),
            "title": point.payload.get("title", ""),
            "file_name": point.payload.get("file_name", ""),
            "file_path": point.payload.get("file_path", "")
        } for point in truncated_root_points]

    async def _get_root_points(self, 
                               collection_name: str, 
                               root_doc_ids: set) -> Dict[int, str]:
        """批量获取顶级父块内容"""
        if not root_doc_ids:
            return {}
        
        root_points = await asyncio.to_thread(
            self.client.retrieve,
            collection_name=collection_name,
            ids=list(root_doc_ids),
            with_payload=True
        )
        
        return root_points
    
    def ping(self) -> bool:
        try:
            return self.client.get_collections().status == "ok"
        except Exception:
            return False